{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 基于LeNet-5的水果识别\n",
    "## 导入包"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import DataLoader, TensorDataset, random_split\n",
    "from torch.utils.data import Dataset\n",
    "from torchvision.transforms import transforms\n",
    "import os\n",
    "import pickle"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 构建模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LeNet5(nn.Module):     \n",
    "    def __init__(self):         \n",
    "        super(LeNet5, self).__init__()         \n",
    "        self.conv1 = nn.Sequential(nn.Conv2d(1, 6, kernel_size=5, stride=1, padding=0),           \n",
    "                nn.ReLU(), \n",
    "                nn.MaxPool2d(kernel_size=2, stride=2, padding=0))         \n",
    "        self.conv2 = nn.Sequential(nn.Conv2d(6, 16, kernel_size=5, stride=1, padding=0),nn.ReLU(),\n",
    "                nn.MaxPool2d(kernel_size=2, stride=2, padding=0))         \n",
    "        self.fc1 = nn.Sequential(nn.Linear(16*5*5, 120),\n",
    "                                 nn.ReLU())         \n",
    "        self.fc2 = nn.Sequential(nn.Linear(120, 84),\n",
    "                                 nn.ReLU())         \n",
    "        self.fc3 = nn.Linear(84, 10)\n",
    "    def forward(self, x):    \n",
    "        x = self.conv1(x)\n",
    "        x = self.conv2(x)\n",
    "        #全连接层均使用的nn.Linear()\n",
    "        #线性结构，输入输出维度均为一维，故需要把数据拉为一维\t\n",
    "        x = x.view(x.size(0), -1)        \n",
    "        x = self.fc1(x)        \n",
    "        x = self.fc2(x)        \n",
    "        x = self.fc3(x)        \n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 构造数据集类"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyDataset(Dataset):\n",
    "    def __init__(self, data_dir, transform=None):\n",
    "        self.data_dir = data_dir\n",
    "        self.transform = transform\n",
    "        self.images = []\n",
    "        # 遍历目录，将所有图片文件路径保存到 self.images 中\n",
    "        for file in os.listdir(self.data_dir):\n",
    "            if file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):\n",
    "                self.images.append(file)\n",
    "                \n",
    "    def __getitem__(self, index):\n",
    "        img_filename = self.images[index]\n",
    "        # 假设图片文件名的首字符为类别标签，如 \"0_apple.png\"\n",
    "        label = int(img_filename[0])\n",
    "        img_path = os.path.join(self.data_dir, img_filename)\n",
    "        img = Image.open(img_path).convert('L')\n",
    "        if self.transform:\n",
    "            img = self.transform(img)\n",
    "        return img, label\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.images)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 构造训练集类"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyTrain():\n",
    "    def __init__(self, dataloader, testDataloader):\n",
    "        # define loss\n",
    "        self.dataloader = dataloader\n",
    "        self.testDataloader = testDataloader\n",
    "        self.device = torch.device(\n",
    "            \"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "        self.net = LeNet5().to(self.device)\n",
    "        self.loss_fuc = nn.CrossEntropyLoss()  # 选择交叉熵损失函数\n",
    "        self.optimizer = torch.optim.SGD(\n",
    "            self.net.parameters(), lr=0.01, momentum=0.9)\n",
    "    def start(self):\n",
    "        EPOCH = 40\n",
    "        for epoch in range(EPOCH):\n",
    "            sum_loss = 0\n",
    "            for i, data in enumerate(self.dataloader):\n",
    "                inputs = data[0]\n",
    "                labels = data[1]\n",
    "                inputs, labels = inputs.to(self.device), labels.to(self.device)\n",
    "                # 梯度清零\n",
    "                self.optimizer.zero_grad()\n",
    "                # 正向传播\n",
    "                output = self.net(inputs)\n",
    "                # 计算误差\n",
    "                loss = self.loss_fuc(output, labels)\n",
    "                # 反向传播\n",
    "                loss.backward()\n",
    "                # 参数更新\n",
    "                self.optimizer.step()\n",
    "                # 打印loss\n",
    "            sum_loss += loss.item()\n",
    "            print('[Epoch:%d, batch:%d] train loss: %.05f' %\n",
    "                  (epoch + 1, i + 1, sum_loss / 100))\n",
    "            sum_loss = 0.0\n",
    "    def predict(self):\n",
    "        print(\"TestResult:\")\n",
    "        for i, data in enumerate(self.testDataloader):\n",
    "            inputs = data[0]\n",
    "            labels = data[1]\n",
    "            # print(labels)\n",
    "            inputs, labels = inputs.to(self.device), labels.to(self.device)\n",
    "            _, predictoutput = torch.max(self.net(inputs), 1)\n",
    "            _, predictoutput = torch.max(self.net(inputs), 1)\n",
    "        print('第%s个预测, 所属类别是%s, 预测结果是%s' % (i+1, labels.item(), predictoutput.item()))\n",
    "\n",
    "    def save(self):\n",
    "        if not os.path.exists('./output'):\n",
    "            os.makedirs('./output')\n",
    "        torch.save(self.net.state_dict(), './output/LeNet5.pth')\n",
    "\n",
    "    def load(self):\n",
    "        if os.path.exists('./output/model.mdl'):\n",
    "            self.net.load_state_dict(torch.load('./output/LeNet5.pth'))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 模型训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Epoch:1, batch:23] train loss: 0.02298\n",
      "[Epoch:2, batch:23] train loss: 0.02281\n",
      "[Epoch:3, batch:23] train loss: 0.02261\n",
      "[Epoch:4, batch:23] train loss: 0.02247\n",
      "[Epoch:5, batch:23] train loss: 0.02137\n",
      "[Epoch:6, batch:23] train loss: 0.02034\n",
      "[Epoch:7, batch:23] train loss: 0.02151\n",
      "[Epoch:8, batch:23] train loss: 0.01936\n",
      "[Epoch:9, batch:23] train loss: 0.01433\n",
      "[Epoch:10, batch:23] train loss: 0.01341\n",
      "[Epoch:11, batch:23] train loss: 0.01589\n",
      "[Epoch:12, batch:23] train loss: 0.01510\n",
      "[Epoch:13, batch:23] train loss: 0.01474\n",
      "[Epoch:14, batch:23] train loss: 0.01075\n",
      "[Epoch:15, batch:23] train loss: 0.01473\n",
      "[Epoch:16, batch:23] train loss: 0.01383\n",
      "[Epoch:17, batch:23] train loss: 0.00837\n",
      "[Epoch:18, batch:23] train loss: 0.01093\n",
      "[Epoch:19, batch:23] train loss: 0.01202\n",
      "[Epoch:20, batch:23] train loss: 0.01157\n",
      "[Epoch:21, batch:23] train loss: 0.01020\n",
      "[Epoch:22, batch:23] train loss: 0.00767\n",
      "[Epoch:23, batch:23] train loss: 0.01218\n",
      "[Epoch:24, batch:23] train loss: 0.00725\n",
      "[Epoch:25, batch:23] train loss: 0.00525\n",
      "[Epoch:26, batch:23] train loss: 0.00468\n",
      "[Epoch:27, batch:23] train loss: 0.00379\n",
      "[Epoch:28, batch:23] train loss: 0.00529\n",
      "[Epoch:29, batch:23] train loss: 0.00400\n",
      "[Epoch:30, batch:23] train loss: 0.00428\n",
      "[Epoch:31, batch:23] train loss: 0.00733\n",
      "[Epoch:32, batch:23] train loss: 0.00333\n",
      "[Epoch:33, batch:23] train loss: 0.00107\n",
      "[Epoch:34, batch:23] train loss: 0.00089\n",
      "[Epoch:35, batch:23] train loss: 0.00082\n",
      "[Epoch:36, batch:23] train loss: 0.00036\n",
      "[Epoch:37, batch:23] train loss: 0.00736\n",
      "[Epoch:38, batch:23] train loss: 0.00153\n",
      "[Epoch:39, batch:23] train loss: 0.00133\n",
      "[Epoch:40, batch:23] train loss: 0.00192\n",
      "TestResult:\n",
      "第182个预测, 所属类别是7, 预测结果是7\n"
     ]
    }
   ],
   "source": [
    "\n",
    "if __name__ == \"__main__\":\n",
    "    transform = transforms.Compose([\n",
    "        transforms.Resize((32, 32)),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(mean=[0.5], std=[0.5])\n",
    "    ])\n",
    "    fullDataset = MyDataset(\"../data/famnist/famnist-all\", transform=transform)\n",
    "    # 按 80%/20% 划分训练集和测试集\n",
    "    trainSize = int(0.8 * len(fullDataset))\n",
    "    testSize = len(fullDataset) - trainSize\n",
    "    trainDataset, testDataset = random_split(fullDataset, [trainSize, testSize])\n",
    "\n",
    "    trainDataloader = DataLoader(trainDataset, batch_size=32, shuffle=True)\n",
    "    testDataloader = DataLoader(testDataset, batch_size=1)\n",
    "    mytrain = MyTrain(trainDataloader, testDataloader)\n",
    "    #mytrain.load()\n",
    "    mytrain.start()\n",
    "    mytrain.predict()\n",
    "    mytrain.save()\t"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
